package com.crawler.search.utils;

import lombok.SneakyThrows;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.bulk.*;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.ByteSizeUnit;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.unit.TimeValue;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.util.function.BiConsumer;


/**
 * 批量插入
 *
 * @author zhangyonghui
 * @Class Name ESBulkProcessor
 * @Create In 下午6:17:12
 */
@Component
public class ESBulkProcessor {

    protected static Logger logger = LoggerFactory.getLogger(ESBulkProcessor.class);

    @Autowired
    public RestHighLevelClient client;

    public static BulkProcessor bulkProcessor;

    @PostConstruct
    public void init() {
        BiConsumer<BulkRequest, ActionListener<BulkResponse>> bulkConsumer =
                (request, bulkListener) -> client.bulkAsync(request, RequestOptions.DEFAULT, bulkListener);
        bulkProcessor = BulkProcessor.builder(bulkConsumer, new BulkProcessor.Listener() {
            @Override
            public void beforeBulk(long executionId, BulkRequest request) {
                logger.info("---插入{}条数据---", request.numberOfActions());
            }

            @SneakyThrows
            @Override
            public void afterBulk(long executionId, BulkRequest request, BulkResponse response) {
                logger.info("---插入{}条数据成功---",
                        request.numberOfActions());
                boolean isFail = response.hasFailures();
                logger.info("请求成功 +++ " + request.numberOfActions() +
                        "\n++++   是否包含错误 " + isFail);
                if (isFail) {
                    BulkItemResponse[] responses = response.getItems();
                    for (BulkItemResponse aResponse : responses) {
                        if (aResponse.isFailed()) {
                            logger.error(" ES 插入错误：" + aResponse.getFailureMessage());
                        }
                    }
                }
            }

            @Override
            public void afterBulk(long executionId, BulkRequest request, Throwable failure) {
                //写入失败后
                logger.error("[es错误]---插入数据失败---", failure);
            }
        })
                .setBulkActions(10000)
                .setBulkSize(new ByteSizeValue(5, ByteSizeUnit.MB))
                .setFlushInterval(TimeValue.timeValueSeconds(5))
                .setConcurrentRequests(15)
                .setBackoffPolicy(
                        BackoffPolicy.exponentialBackoff(TimeValue.timeValueMillis(100), 3))
                .build();
    }

    public void addIndexRequest(IndexRequest indexRequest){
        bulkProcessor.add(indexRequest);
    }


}
